import xgboost as xgb
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_curve
import numpy as np

def read_data():
    data = pd.read_csv('D:/lung_cancer/data/data.csv')
    features = []
    labels = []
    for i in range(len(data)):
        if data['cancer_type'][i]-1 == 0 or data['cancer_type'][i]-1 == 1:
            one_feature = [data['z'][i], data['x'][i], data['y'][i], data['r'][i],
                           data['patientWeight'][i], data['patientSex'][i], data['patientAge'][i],
                           data['part_suvmax'][i], data['suv_avg'][i], data['suv_std'][i]]
            features.append(one_feature)
            labels.append(data['cancer_type'][i]-1)
    print(features)
    print(len(features))
    print(labels)
    print(len(labels))

    # 只用第一类和第二类数据

    X_train, X_test, y_train, y_test = train_test_split(np.array(features, dtype=np.float),
                                                        np.array(labels, dtype=np.int),
                                                        test_size=0.2, random_state=1234565)

    print('Xtrain type: ', type(X_train))
    print('y_train type: ', type(y_train))
    return X_train, X_test, y_train, y_test

def train():
    train_features, test_features, train_labels, test_labels = read_data()
    dtrain = xgb.DMatrix(train_features, label=train_labels)
    dtest = xgb.DMatrix(test_features, label=test_labels)

    # dtrain = xgb.DMatrix('D:/lung_cancer/xgboost-master/xgboost-master/demo/data/agaricus.txt.train')
    # dtest = xgb.DMatrix('D:/lung_cancer/xgboost-master/xgboost-master/demo/data/agaricus.txt.test')

    param = {
        'booster': 'dart',
        'max_depth': 5,
        'learning_rate': 0.1,
        'objective': 'binary:logistic',
        'sample_type': 'uniform',
        'normalize_type': 'tree',
        'rate_drop': 0.1,
        'skip_drop': 0.5
    }

    num_round = 50
    bst = xgb.train(param, dtrain, num_round)
    preds = bst.predict(dtest)
    print(preds)

if __name__ == '__main__':
    train()